package ca

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"math/big"
	"os"
	"time"
)

type CA_Info struct {
	Country            string
	Organization       string
	OrganizationalUnit string
	Province           string
	CommonName         string
	Locality           string
}

func LoadCAcert(caPath string) *x509.Certificate {
	caFile, err := ioutil.ReadFile(caPath)
	if err != nil {
		panic(err)
	}
	caBlock, _ := pem.Decode(caFile)
	cert, err := x509.ParseCertificate(caBlock.Bytes) // cert
	if err != nil {
		panic(err)
	}
	return cert
}

func loadCAkey(caPath string) *rsa.PrivateKey {
	keyFile, err := ioutil.ReadFile(caPath)
	if err != nil {
		panic(err)
	}
	keyBlock, _ := pem.Decode(keyFile)
	praKey, err := x509.ParsePKCS1PrivateKey(keyBlock.Bytes) // praKey
	if err != nil {
		panic(err)
	}
	return praKey
}

// 根据rootDir中存放的CA证书生成一个新的证书
func GenerateCert(ca_info CA_Info, saveDir string, rootDir string) {

	rootCAcert := LoadCAcert(rootDir + "/cert.crt")
	rootCAkey := loadCAkey(rootDir + "/key.pem")

	template := &x509.Certificate{
		SerialNumber: big.NewInt(1), //证书序列号
		Subject: pkix.Name{
			Country:            []string{ca_info.Country},
			Organization:       []string{ca_info.Organization},
			OrganizationalUnit: []string{ca_info.OrganizationalUnit},
			Province:           []string{ca_info.Province},
			CommonName:         "localhost",
			Locality:           []string{ca_info.Locality},
		},
		NotBefore:             time.Now(),                                                                 //证书有效期开始时间
		NotAfter:              time.Now().AddDate(1, 0, 0),                                                //证书有效期结束时间
		BasicConstraintsValid: true,                                                                       //基本的有效性约束
		IsCA:                  false,                                                                      //是否是根证书
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth}, //证书用途(客户端认证，数据加密)
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageDataEncipherment,
		//EmailAddresses: []string{"test@test.com"},
		//IPAddresses:    []net.IP{net.ParseIP("192.168.1.59")},
	}

	priKey, err := rsa.GenerateKey(rand.Reader, 2048)

	// 由CA证书作为根证书，生成一个新的证书，使用的是CA的私钥rootCAkey的进行签名，把证书持有者的公钥priKey.PublicKey记录到证书里
	derBytes, err := x509.CreateCertificate(rand.Reader, template, rootCAcert, &priKey.PublicKey, rootCAkey)
	if err != nil {
		fmt.Println("Failed to create certificate:", err)
		return
	}

	// 将证书保存到文件
	certOut, err := os.Create(saveDir + "/cert.crt")
	if err != nil {
		fmt.Println("Failed to create certificate file:", err)
		return
	}
	defer certOut.Close()
	pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})

	// 将私钥保存到文件
	keyOut, err := os.OpenFile(saveDir+"/key.pem", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		fmt.Println("Failed to open key.pem for writing:", err)
		return
	}
	defer keyOut.Close()
	pem.Encode(keyOut, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priKey)})
}

// 验证证书是否由CA签发
func Verify(cert *x509.Certificate, rootPath string) bool {
	rootCert := LoadCAcert(rootPath)

	roots := x509.NewCertPool()
	roots.AddCert(rootCert)
	opts := x509.VerifyOptions{
		Roots: roots,
	}

	chain, err := cert.Verify(opts)
	if err != nil {
		return false
	} else {
		return rootCert.Equal(chain[0][len(chain[0])-1])
	}
}
